package tools

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"

	"mcp-db-gateway/pkg/db"
)

type DBExecutor struct {
	dbManager *db.Manager
}

func NewDBExecutor(dbManager *db.Manager) *DBExecutor {
	return &DBExecutor{
		dbManager: dbManager,
	}
}

type ExecuteSQLParams struct {
	DBName    string `json:"db_name"`
	SQL       string `json:"sql"`
	SessionID string `json:"session_id,omitempty"`
}

func (e *DBExecutor) ExecuteSQL(ctx context.Context, paramsJSON string) (interface{}, error) {
	var params ExecuteSQLParams
	if err := json.Unmarshal([]byte(paramsJSON), &params); err != nil {
		return nil, fmt.Errorf("invalid parameters: %w", err)
	}

	trimmedSQL := strings.TrimSpace(strings.ToUpper(params.SQL))

	// Handle transaction control statements
	switch trimmedSQL {
	case "BEGIN", "START TRANSACTION":
		sessionID, err := e.dbManager.BeginTx(params.DBName)
		if err != nil {
			return nil, err
		}
		return map[string]string{"session_id": sessionID}, nil
	case "COMMIT":
		if params.SessionID == "" {
			return nil, fmt.Errorf("session_id is required for COMMIT")
		}
		return nil, e.dbManager.CommitTx(params.SessionID)
	case "ROLLBACK":
		if params.SessionID == "" {
			return nil, fmt.Errorf("session_id is required for ROLLBACK")
		}
		return nil, e.dbManager.RollbackTx(params.SessionID)
	}

	// If session_id is provided, execute within a transaction
	if params.SessionID != "" {
		tx, ok := e.dbManager.GetTx(params.SessionID)
		if !ok {
			return nil, fmt.Errorf("transaction with session ID '%s' not found or has expired", params.SessionID)
		}
		if strings.HasPrefix(trimmedSQL, "SELECT") {
			return e.txQuery(ctx, tx, params.SQL)
		}
		return e.txExec(ctx, tx, params.SQL)
	}

	// Fallback to non-transactional execution
	database, ok := e.dbManager.GetDB(params.DBName)
	if !ok {
		return nil, fmt.Errorf("database alias '%s' not found", params.DBName)
	}

	if strings.HasPrefix(trimmedSQL, "SELECT") {
		return e.handleQuery(ctx, database, params.SQL)
	}

	return e.handleExec(ctx, database, params.SQL)
}

func (e *DBExecutor) handleQuery(ctx context.Context, db *sql.DB, query string) (interface{}, error) {
	rows, err := db.QueryContext(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	var results []map[string]interface{}
	for rows.Next() {
		values := make([]interface{}, len(columns))
		valuePtrs := make([]interface{}, len(columns))
		for i := range columns {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return nil, err
		}

		rowData := make(map[string]interface{})
		for i, col := range columns {
			val := values[i]
			// Handle potential byte slices and convert them to strings for JSON serialization
			if b, ok := val.([]byte); ok {
				rowData[col] = string(b)
			} else {
				rowData[col] = val
			}
		}
		results = append(results, rowData)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return results, nil
}

func (e *DBExecutor) txQuery(ctx context.Context, tx *sql.Tx, query string) (interface{}, error) {
	rows, err := tx.QueryContext(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	var results []map[string]interface{}
	for rows.Next() {
		values := make([]interface{}, len(columns))
		valuePtrs := make([]interface{}, len(columns))
		for i := range columns {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return nil, err
		}

		rowData := make(map[string]interface{})
		for i, col := range columns {
			val := values[i]
			if b, ok := val.([]byte); ok {
				rowData[col] = string(b)
			} else {
				rowData[col] = val
			}
		}
		results = append(results, rowData)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return results, nil
}

func (e *DBExecutor) handleExec(ctx context.Context, db *sql.DB, query string) (interface{}, error) {
	result, err := db.ExecContext(ctx, query)
	if err != nil {
		return nil, err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return nil, fmt.Errorf("failed to get rows affected: %w", err)
	}

	return map[string]interface{}{
		"rows_affected": rowsAffected,
	}, nil
}

func (e *DBExecutor) txExec(ctx context.Context, tx *sql.Tx, query string) (interface{}, error) {
	result, err := tx.ExecContext(ctx, query)
	if err != nil {
		return nil, err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return nil, fmt.Errorf("failed to get rows affected: %w", err)
	}

	return map[string]interface{}{
		"rows_affected": rowsAffected,
	}, nil
}

func (e *DBExecutor) GetDBAliases() []string {
	return e.dbManager.GetDBAliases()
} 